import time
import serial
from PIL import Image
import argparse
import matplotlib.pyplot as plt

def send_image(ser, path):
    img = Image.open(path).convert('L')
    pixels = list(img.getdata())
    if len(pixels)!=28*28:
        raise ValueError("图像尺寸必须 28x28")
    # 每次发送前清理串口缓存
    ser.reset_input_buffer()
    ser.reset_output_buffer()
    # 发送数据并读取结果
    ser.write(b"START")
    # 分块发送像素，避免串口缓冲区溢出
    chunk_size = 100
    for i in range(0, len(pixels), chunk_size):
        ser.write(bytes(pixels[i:i+chunk_size]))
        ser.flush()
        time.sleep(0.01)
    ser.write(b"END")
    ser.flush()
    # 等待 FPGA 处理
    time.sleep(0.2)
    # 直接读取 DONE（4 字节）
    ack = ser.read(4)
    if ack != b"DONE":
        raise ValueError(f"未收到 DONE，返回: {ack}")
    # 读取标签字节
    label_byte = ser.read(1)
    try:
        label = label_byte.decode()
    except:
        label = str(label_byte)
    return img, label

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="发送图像到 FPGA 的串口接口")
    parser.add_argument("--port", default="COM12", help="串口端口，如 COM12")
    parser.add_argument("--baudrate", type=int, default=115200, help="波特率，如 115200")
    args = parser.parse_args()
    # 打开并初始化串口
    ser = serial.Serial(args.port, args.baudrate, timeout=5)
    ser.reset_input_buffer()
    ser.reset_output_buffer()
    time.sleep(2)
    try:
        while True:
            path = input("请输入图像文件路径（输入 q 退出）：")
            
            if path.lower() == 'q':
                break
            try:
                img, label = send_image(ser, path)
                line = ser.readline()
                if line: print(line)
                plt.figure()
                plt.imshow(img, cmap='gray')
                plt.title(f"Identification result: {label}")
                plt.axis('off')
                plt.show()
            except Exception as e:
                print(f"Error: {e}")
    finally:
        ser.close()